import os
import random

from PIL import Image
from torchvision import transforms


def process_dataset_image_name(dataset_path):
    """
    将数据集中的图片文件转换为按顺序命名
    :param dataset_path: 人脸数据集路径
    :return: None
    """
    print("1.将数据集中的图片文件转换为按顺序命名")
    # 遍历每一个分类
    for dir_name in os.listdir(dataset_path):
        image_dir_path = os.path.join(dataset_path, dir_name)
        if os.path.isdir(image_dir_path):  # 判断是否为文件夹
            # 遍历每个分类文件夹下的图片并rename
            for index, image_name in enumerate(os.listdir(image_dir_path)):
                image_path = os.path.join(image_dir_path, image_name)
                new_image_path = os.path.join(image_dir_path, f"{index + 1:03d}.jpg")
                os.rename(image_path, new_image_path)


def create_annotation(dataset_path):
    """
    遍历数据集文件夹生成标注文件
    :param dataset_path: 人脸数据集路径
    :return: 标注文件的路径 annotation_path
    """
    print("2.遍历数据集文件夹生成标注文件")
    cwd = os.getcwd()
    annotation_path = os.path.join(dataset_path, "cls_train.txt")
    with open(annotation_path, 'w') as fp:
        # 遍历每一个分类
        for cls_id, cls_name in enumerate(os.listdir(dataset_path)):
            photos_path = os.path.join(dataset_path, cls_name)
            if os.path.isdir(photos_path):  # 判断是否为文件夹
                # 遍历每个分类文件夹下的图片并写入标注文件
                for photo_name in os.listdir(photos_path):
                    fp.write(f"{cls_id - 1};{cwd}/{os.path.join(photos_path, photo_name)}")
                    fp.write('\n')
    return annotation_path


def create_pair(dataset_path):
    """
    制作 pair 文件
    :param dataset_path: 人脸数据集路径
    :return: None
    """
    print("3.制作 pair 文件")
    name_list = ["jiangwen", "pengyuyan", "zhangziyi"]
    images_size = [103, 114, 100]
    pair_path = os.path.join(dataset_path, "pairs.txt")
    with open(pair_path, 'a') as fp:
        for _ in range(10):
            for _ in range(100):
                index = random.randint(0, 2)
                image_index_1 = random.randint(1, images_size[index])
                image_index_2 = random.randint(1, images_size[index])
                fp.write(f"{name_list[index]}\t{image_index_1}\t{image_index_2}\n")
            for _ in range(100):
                index1 = random.randint(0, 2)
                image_index_1 = random.randint(1, images_size[index1])
                index2 = (index1 + 1) % 2
                image_index_2 = random.randint(1, images_size[index2])
                fp.write(f"{name_list[index1]}\t{image_index_1}\t{name_list[index2]}\t{image_index_2}\n")


def letterbox_image(image, size):
    image = image.convert("RGB")
    iw, ih = image.size
    w, h = size
    scale = min(w / iw, h / ih)
    nw = int(iw * scale)
    nh = int(ih * scale)

    image = image.resize((nw, nh), Image.BICUBIC)
    new_image = Image.new('RGB', size, (128, 128, 128))
    new_image.paste(image, ((w - nw) // 2, (h - nh) // 2))
    return new_image


def horizontal_flip(image):
    HF = transforms.RandomHorizontalFlip()
    hf_image = HF(image)
    return hf_image


def vertical_flip(image):
    VF = transforms.RandomVerticalFlip()
    vf_image = VF(image)
    return vf_image


def random_rotation(image):
    RR = transforms.RandomRotation(degrees=(10, 80))
    rr_image = RR(image)
    return rr_image


def process_image(image, size):
    image = letterbox_image(image, size)
    if random.random() < 0.5:
        image = horizontal_flip(image)
    if random.random() < 0.5:
        image = vertical_flip(image)
    if random.random() < 0.5:
        image = random_rotation(image)
    return image


if __name__ == '__main__':
    process_dataset_image_name("images/face")
    create_annotation("images/face")
    create_pair("images/face")
